import pickle
from absl import app
from absl import flags
import matplotlib

matplotlib.use('Agg')  # 无图形界面模式
import matplotlib.pyplot as plt
import numpy as np
import os
from pyevtk.hl import pointsToVTK
from matplotlib import animation
import csv
from datetime import datetime
import sys

# 新增Lacey指数相关参数
flags.DEFINE_bool("compute_lacey", True, help="Whether to compute Lacey index")
flags.DEFINE_string("type1", "6", help="First particle type for mixing (comma separated)")
flags.DEFINE_string("type2", "0", help="Second particle type for mixing")
flags.DEFINE_float("mass_cut_off", 1e-12, help="Minimum total mass in a bin to be considered")
flags.DEFINE_string("bins", "2,5,5", help="Number of bins in x,y,z directions")
# 原始参数
flags.DEFINE_string("rollout_dir", None, help="Directory where rollout.pkl are located")
flags.DEFINE_string("rollout_name", None, help="Name of rollout `.pkl` file")
flags.DEFINE_integer("step_stride", 1, help="Stride of steps to skip.")
flags.DEFINE_bool("change_yz", False, help="Change y and z axis.")
flags.DEFINE_enum("output_mode", "gif", ["gif", "vtk"], help="Type of render output")

FLAGS = flags.FLAGS

TYPE_TO_COLOR = {
    1: "red",  # for droplet
    3: "black",  # Boundary particles.
    0: "green",  # Rigid solids.
    7: "magenta",  # Goop.
    6: "gold",  # Sand.
    5: "blue",  # Water.
}


# ==== 新增 Lacey 指数计算功能 ====
def grid(minCoords, maxCoords, bins):
    """创建网格点坐标"""
    lengths = maxCoords - minCoords
    division_size = lengths / bins
    x_coords = np.arange(minCoords[0] + (division_size[0] / 2), maxCoords[0], division_size[0])
    y_coords = np.arange(minCoords[1] + (division_size[1] / 2), maxCoords[1], division_size[1])
    z_coords = np.arange(minCoords[2] + (division_size[2] / 2), maxCoords[2], division_size[2])
    return np.array([[x, y, z] for x in x_coords for y in y_coords for z in z_coords]), division_size


def bin_particles(b_coords, div_size, particles, mass_1, mass_2, conc):
    """将粒子分配到网格中并计算质量分数"""
    p_type_1 = particles[0, 4]

    for i, coord in enumerate(b_coords):
        mins = coord - div_size / 2
        maxs = coord + div_size / 2
        in_bin = (particles[:, 0] > mins[0]) & (particles[:, 0] < maxs[0]) & \
                 (particles[:, 1] > mins[1]) & (particles[:, 1] < maxs[1]) & \
                 (particles[:, 2] > mins[2]) & (particles[:, 2] < maxs[2])

        type_mask = particles[:, 4] == p_type_1
        mass_1[i] = np.sum(particles[in_bin & type_mask, 3])
        mass_2[i] = np.sum(particles[in_bin & ~type_mask, 3])

        total = mass_1[i] + mass_2[i]
        if total > 0:
            conc[i] = mass_1[i] / total


def calculate_lacey_index(mass_1, mass_2, conc, cut_off):
    """计算Lacey混合指数 - 防止负值"""
    total_mass1 = np.sum(mass_1)
    total_mass2 = np.sum(mass_2)
    total_mass = total_mass1 + total_mass2

    if total_mass < 1e-12:  # 防止除以零
        return 0.0

    P = total_mass1 / total_mass
    S0 = P * (1 - P)

    if S0 < 1e-12:  # P接近0或1的情况
        return 0.0

    mass_total = mass_1 + mass_2
    valid_bins = mass_total > cut_off
    n_bins = np.sum(valid_bins)

    if n_bins < 2:  # 至少需要2个有效bin
        return 0.0

    # 计算有效网格中的平均方差
    S = np.mean((conc[valid_bins] - P) ** 2)

    # 防止数值错误
    if S < 0:
        S = 0

    # 计算随机混合的方差
    Sr = S0 / max(1, n_bins)  # 防止除以零

    # 确保分母为正
    denominator = S0 - Sr
    if denominator < 0:
        return 0.0

    # 计算并限制在[0,1]范围内
    M = (S0 - S) / denominator
    return max(0.0, min(1.0, M))  # 确保在0-1范围内


# ==== 修改后的Render类 ====
class Render():
    def __init__(self, input_dir, input_name):
        # 解析Lacey参数
        self.compute_lacey = FLAGS.compute_lacey
        self.type1_set = set(map(int, FLAGS.type1.split(',')))
        self.type2_set = set(map(int, FLAGS.type2.split(',')))
        self.bins = list(map(int, FLAGS.bins.split(',')))
        self.mass_cut_off = FLAGS.mass_cut_off

        self.input_dir = input_dir
        self.input_name = input_name
        self.output_dir = input_dir
        self.output_name = input_name

        with open(f"{self.input_dir}{self.input_name}.pkl", "rb") as file:
            rollout_data = pickle.load(file)
        self.rollout_data = rollout_data

        # 初始化质量映射 - 假设所有粒子质量为1
        self.particle_masses = np.ones(len(rollout_data["particle_types"]))

        rollout_cases = [
            ["ground_truth_rollout", "Reality"],
            ["predicted_rollout", "GNS"]
        ]
        self.rollout_cases = rollout_cases

        trajectory = {}
        for rollout_case in rollout_cases:
            trajectory[rollout_case[0]] = np.concatenate(
                [rollout_data["initial_positions"], rollout_data[rollout_case[0]]], axis=0
            )
        self.trajectory = trajectory
        self.loss = self.rollout_data['loss'].item() if 'loss' in self.rollout_data else 0.0

        # 轨迹信息
        self.dims = trajectory[rollout_cases[0][0]].shape[2]
        self.num_particles = trajectory[rollout_cases[0][0]].shape[1]
        self.num_steps = trajectory[rollout_cases[0][0]].shape[0]
        self.boundaries = rollout_data["metadata"]["bounds"]
        self.particle_type = rollout_data["particle_types"]

        # Lacey指数存储
        self.lacey_indices = {"ground_truth": [], "predicted": []}
        self.lacey_times = []

    def color_map(self):
        color_map = np.empty(self.num_particles, dtype="object")
        for material_id, color in TYPE_TO_COLOR.items():
            color_index = np.where(np.array(self.particle_type) == material_id)
            color_map[color_index] = color
        return list(color_map)

    def color_mask(self):
        color_mask = []
        for material_id, color in TYPE_TO_COLOR.items():
            mask = np.array(self.particle_type) == material_id
            if mask.any():
                color_mask.append([mask, color])
        return color_mask

    def prepare_lacey_data(self, positions, step):
        """准备Lacey指数计算所需的数据 - 重构以提高复用性"""
        if not self.compute_lacey:
            return None, None, None

        # 1. 筛选有效粒子
        combined_types = list(self.type1_set) + list(self.type2_set)
        type_mask = np.isin(self.particle_type, combined_types)
        valid_indices = np.where(type_mask)[0]

        if len(valid_indices) == 0:  # 没有有效粒子
            print(f"Step {step}: No valid particles found for types {combined_types}")
            return None, None, None

        valid_particles = positions[valid_indices]
        valid_types = self.particle_type[valid_indices]
        valid_masses = self.particle_masses[valid_indices]

        # 2. 创建类型标记
        type_flags = np.where(
            np.isin(valid_types, list(self.type1_set)),
            1.0,  # 第一类粒子
            0.0  # 第二类粒子
        )

        # 3. 创建带质量的数据结构
        particles_with_type_mass = np.column_stack((
            valid_particles,
            valid_masses,
            type_flags
        ))

        # # 4. 获取边界 - 基于实际粒子位置动态调整
        # minCoords = np.min(valid_particles, axis=0) - 0.0005
        # maxCoords = np.max(valid_particles, axis=0) + 0.0005

        # 使用固定边界
        minCoords = np.array([0.0, 0.0, 0.0])
        maxCoords = np.array([0.03, 0.11, 0.11])

        bins_3d = self.bins

        # 5. 创建网格
        grid_coords, div_size = grid(minCoords, maxCoords, bins_3d)

        # 6. 初始化数组
        mass_1 = np.zeros(len(grid_coords))
        mass_2 = np.zeros(len(grid_coords))
        conc = np.zeros(len(grid_coords))

        # 7. 分配粒子到网格
        bin_particles(grid_coords, div_size, particles_with_type_mass, mass_1, mass_2, conc)

        # 调试信息：统计有效网格数量
        total_mass = mass_1 + mass_2
        valid_bins = total_mass > self.mass_cut_off
        print(f"Step {step}: Valid bins: {np.sum(valid_bins)}/{len(grid_coords)}")

        return mass_1, mass_2, conc

    def calculate_and_store_lacey(self, positions, case_name, step):
        """计算并存储Lacey指数 - 使用prepare_lacey_data"""
        if not self.compute_lacey:
            return

        # 调用prepare_lacey_data函数准备数据
        mass_1, mass_2, conc = self.prepare_lacey_data(positions, step)

        # 如果没有有效数据，则跳过计算
        if mass_1 is None or mass_2 is None or conc is None:
            print(f"Step {step}: prepare_lacey_data 返回了空值，跳过计算")
            return

        # 计算Lacey指数
        lacey_index = calculate_lacey_index(mass_1, mass_2, conc, self.mass_cut_off)

        # 存储结果
        if "ground" in case_name.lower() or "reality" in case_name.lower():
            self.lacey_indices["ground_truth"].append(lacey_index)
        else:
            self.lacey_indices["predicted"].append(lacey_index)

        print(f"Step {step}: Lacey Index ({case_name}) = {lacey_index:.4f}")

    def save_lacey_results(self):
        """保存Lacey指数结果 - 修正确保正确显示曲线和输出数据"""
        if not self.compute_lacey:
            return

        # 创建输出目录
        output_dir = os.path.join(self.output_dir, "lacey_results")
        os.makedirs(output_dir, exist_ok=True)

        # 确保我们有足够的数据
        if not self.lacey_indices["ground_truth"] and not self.lacey_indices["predicted"]:
            print("警告: 没有计算到任何Lacey指数数据，跳过保存结果")
            return

        # 获取两种数据的最小长度
        min_length = min(
            len(self.lacey_indices["ground_truth"]) if self.lacey_indices["ground_truth"] else float('inf'),
            len(self.lacey_indices["predicted"]) if self.lacey_indices["predicted"] else float('inf')
        )

        # 如果没有有效数据，使用默认长度
        if min_length == float('inf'):
            min_length = 0

        # 确保不超过597个步长
        max_step = min(597, min_length) if min_length > 0 else 0

        # 修改点：每隔10步取一次数据点
        stride = 20  # 每隔10个时间步
        sampled_steps = []  # 采样的时间步
        sampled_gt = []  # 采样的真实数据
        sampled_pred = []  # 采样的预测数据

        # 收集每隔10步的数据
        for i in range(0, max_step, stride):
            sampled_steps.append(i)
            if i < len(self.lacey_indices["ground_truth"]):
                sampled_gt.append(self.lacey_indices["ground_truth"][i])
            else:
                sampled_gt.append(None)

            if i < len(self.lacey_indices["predicted"]):
                sampled_pred.append(self.lacey_indices["predicted"][i])
            else:
                sampled_pred.append(None)

        # 如果有有效数据才绘图
        if sampled_steps:
            plt.figure(figsize=(10, 6))

            # 只绘制实际存在的步数
            gt_exists = any(val is not None for val in sampled_gt)
            pred_exists = any(val is not None for val in sampled_pred)

            if gt_exists:
                valid_gt_steps = [step for step, val in zip(sampled_steps, sampled_gt) if val is not None]
                valid_gt_vals = [val for val in sampled_gt if val is not None]
                plt.plot(valid_gt_steps, valid_gt_vals, 'b-', lw=2, marker='o', markersize=5, label='Ground Truth')

            if pred_exists:
                valid_pred_steps = [step for step, val in zip(sampled_steps, sampled_pred) if val is not None]
                valid_pred_vals = [val for val in sampled_pred if val is not None]
                plt.plot(valid_pred_steps, valid_pred_vals, 'r--', lw=2, marker='s', markersize=5, label='Predicted')

            # 设置坐标轴范围
            max_x = sampled_steps[-1]
            plt.xlim(0, max_x)
            plt.ylim(0, 1.0)

            plt.xlabel('Time Step', fontsize=12)
            plt.ylabel('Lacey Mixing Index', fontsize=12)
            plt.title(f'Mixing Index Evolution: {self.output_name}')

            # 根据实际步数设置刻度
            if max_x <= 600:
                plt.xticks(np.arange(0, max_x + 1, 100))
            else:
                plt.xticks(np.arange(0, max_x + 1, max_x // 6))

            plt.yticks(np.arange(0, 1.1, 0.2))

            if gt_exists or pred_exists:
                plt.legend()
            plt.grid(True, alpha=0.3)
            plt.tight_layout()

            # 保存图像
            plt.savefig(os.path.join(output_dir, f"{self.output_name}_lacey_curve.png"), dpi=300)
            plt.close()

            print(f"Lacey曲线图已保存到: {os.path.join(output_dir, f'{self.output_name}_lacey_curve.png')}")
        else:
            print("警告: 没有足够的Lacey指数数据来创建曲线图")

        # 保存CSV报告（分为两列）- 也使用采样数据
        csv_path = os.path.join(output_dir, f"{self.output_name}_lacey_report.csv")
        with open(csv_path, 'w', newline='') as f:
            writer = csv.writer(f)
            writer.writerow(['Step', 'Ground Truth', 'Predicted'])

            # 使用采样数据
            for step, gt_val, pred_val in zip(sampled_steps, sampled_gt, sampled_pred):
                row_gt = round(gt_val, 4) if gt_val is not None else ""
                row_pred = round(pred_val, 4) if pred_val is not None else ""
                writer.writerow([step, row_gt, row_pred])

        print(f"Lacey CSV报告已保存到: {csv_path}")

    def render_gif_animation(self, point_size=1, timestep_stride=3,
                             vertical_camera_angle=20, viewpoint_rotation=0.5,
                             change_yz=False):
        fig = plt.figure()
        if self.dims == 2:
            ax1 = fig.add_subplot(1, 2, 1, projection='rectilinear')
            ax2 = fig.add_subplot(1, 2, 2, projection='rectilinear')
            axes = [ax1, ax2]
        elif self.dims == 3:
            ax1 = fig.add_subplot(1, 2, 1, projection='3d')
            ax2 = fig.add_subplot(1, 2, 2, projection='3d')
            axes = [ax1, ax2]

        trajectory_datacases = [self.rollout_cases[0][0], self.rollout_cases[1][0]]
        render_datacases = [self.rollout_cases[0][1], self.rollout_cases[1][1]]

        xboundary = self.boundaries[0]
        yboundary = self.boundaries[1]
        if self.dims == 3:
            zboundary = self.boundaries[2]

        color_mask = self.color_mask()

        # 存储每个时间步Lacey指数
        lacey_step_indices = []

        def animate(i):
            print(f"Render step {i}/{self.num_steps}")

            fig.clear()
            for j, datacase in enumerate(trajectory_datacases):
                # 计算当前时间步的Lacey指数
                positions = self.trajectory[datacase][i]
                self.calculate_and_store_lacey(positions, render_datacases[j], i)

                if self.dims == 2:
                    axes[j] = fig.add_subplot(1, 2, j + 1)
                    axes[j].set_aspect("equal")
                    axes[j].set_xlim([float(xboundary[0]), float(xboundary[1])])
                    axes[j].set_ylim([float(yboundary[0]), float(yboundary[1])])
                    for mask, color in color_mask:
                        axes[j].scatter(positions[mask, 0], positions[mask, 1],
                                        s=point_size, color=color)
                    axes[j].grid(True, which='both')
                    axes[j].set_title(render_datacases[j])

                elif self.dims == 3:
                    axes[j] = fig.add_subplot(1, 2, j + 1, projection='3d')
                    if not change_yz:
                        axes[j].set_xlim([float(xboundary[0]), float(xboundary[1])])
                        axes[j].set_ylim([float(yboundary[0]), float(yboundary[1])])
                        axes[j].set_zlim([float(zboundary[0]), float(zboundary[1])])
                        for mask, color in color_mask:
                            axes[j].scatter(positions[mask, 0], positions[mask, 1],
                                            positions[mask, 2], s=point_size, color=color)
                        axes[j].set_box_aspect((float(xboundary[1]) - float(xboundary[0]),
                                                float(yboundary[1]) - float(yboundary[0]),
                                                float(zboundary[1]) - float(zboundary[0])))
                        axes[j].view_init(elev=vertical_camera_angle, azim=i * viewpoint_rotation)

                    else:
                        axes[j].set_xlim([float(xboundary[0]), float(xboundary[1])])
                        axes[j].set_ylim([float(zboundary[0]), float(zboundary[1])])
                        axes[j].set_zlim([float(yboundary[0]), float(yboundary[1])])
                        for mask, color in color_mask:
                            axes[j].scatter(positions[mask, 0], positions[mask, 2],
                                            positions[mask, 1], s=point_size, color=color)
                        axes[j].set_box_aspect((float(xboundary[1]) - float(xboundary[0]),
                                                float(zboundary[1]) - float(zboundary[0]),
                                                float(yboundary[1]) - float(yboundary[0])))
                        axes[j].view_init(elev=vertical_camera_angle, azim=i * viewpoint_rotation)

                    axes[j].grid(True, which='both')
                    axes[j].set_title(render_datacases[j])

            fig.suptitle(f"{i}/{self.num_steps}, Total MSE: {self.loss:.2e}")

        ani = animation.FuncAnimation(
            fig, animate, frames=np.arange(0, self.num_steps, timestep_stride), interval=50)

        ani.save(f'{self.output_dir}{self.output_name}.gif', dpi=300, fps=10, writer='imagemagick')
        print(f"Animation saved to: {self.output_dir}{self.output_name}.gif")

        # 保存Lacey结果
        self.save_lacey_results()

    def write_vtk(self):
        for rollout_case, label in self.rollout_cases:
            path = f"{self.output_dir}{self.output_name}_vtk-{label}"
            if not os.path.exists(path):
                os.makedirs(path)

            initial_position = self.trajectory[rollout_case][0]

            for i, coord in enumerate(self.trajectory[rollout_case]):
                # 计算当前时间步的Lacey指数
                self.calculate_and_store_lacey(coord, label, i)

                # 确保位移计算使用副本而不是视图
                disp = np.linalg.norm(coord - initial_position, axis=1)

                # 处理2D情况
                if self.dims == 2:
                    z_data = np.zeros(coord.shape[0], dtype=np.float32)
                else:
                    z_data = coord[:, 2].astype(np.float32)

                # 确保所有数组都是C连续且类型正确
                x_data = np.ascontiguousarray(coord[:, 0].astype(np.float32))
                y_data = np.ascontiguousarray(coord[:, 1].astype(np.float32))
                z_data = np.ascontiguousarray(z_data)
                disp_data = np.ascontiguousarray(disp.astype(np.float32))

                pointsToVTK(
                    f"{path}/points{i}",
                    x_data,
                    y_data,
                    z_data,
                    data={"displacement": disp_data}
                )

        print(f"vtk saved to: {self.output_dir}{self.output_name}...")
        self.save_lacey_results()


# ==== 主函数 ====
def main(_):
    if not FLAGS.rollout_dir:
        raise ValueError("A `rollout_dir` must be passed.")
    if not FLAGS.rollout_name:
        raise ValueError("A `rollout_name`must be passed.")

    render = Render(input_dir=FLAGS.rollout_dir, input_name=FLAGS.rollout_name)

    if FLAGS.output_mode == "gif":
        render.render_gif_animation(
            point_size=1,
            timestep_stride=FLAGS.step_stride,
            vertical_camera_angle=20,
            viewpoint_rotation=0.0,
            change_yz=FLAGS.change_yz
        )
    elif FLAGS.output_mode == "vtk":
        render.write_vtk()


if __name__ == '__main__':
    app.run(main)